package fasttask

import (
    "cloud_worm/util/umodels/mysqldb"
    "cloud_worm/util/umodels/mysqldb/dbtask"
    "cloud_worm/util/umodels/mysqldb/dbdata"
    "cloud_worm/util/umodels/mysqldb/dbmanager"
    "github.com/astaxie/beego/logs"
    "fmt"
    "cloud_worm/util/uconst"
    "cloud_worm/util/uid"
    "cloud_worm/util/proto/urpc"
    "time"
    "net"
    "strings"
    "cloud_worm/wormweb/controllers/ctrsecurity/webtask/slowtask"
    "cloud_worm/wormweb/controllers/ctrinner/taskqueue"
    "cloud_worm/wormweb/controllers/ctrinner/analyzer"
    "cloud_worm/wormweb/controllers/ctrutil"
)

type TaskItem struct {
    NotifyChan      chan int
    TaskId          string                  // 任务ID，task id
    WormId           string                 // worm id，标识一个机器
    WeaponName      string                  // 插件名字
    Input           string                  // 执行一个weapon的input输入信息
}

/**
 @Description：获取task中ok和failed的个数
 @Param:
 @Return：
 */
func GetTaskCount(task *dbtask.WeaponTask) (int, int) {

    wLog := &dbdata.WeaponLog{}
    okCnt, err := wLog.Count(mysqldb.SelectStrAll, mysqldb.SelectIntAll, task.TaskId, ctrutil.WEAPON_LOG_DONE)
    if err != nil {
        logs.Error(fmt.Sprintf("count weapon log done for `%s` error: %s", task.TaskId, err.Error()))
    }
    failedCnt, err := wLog.Count(mysqldb.SelectStrAll, mysqldb.SelectIntAll, task.TaskId, ctrutil.WEAPON_LOG_FAILED)
    if err != nil {
        logs.Error(fmt.Sprintf("count weapon log fail for `%s` error: %s", task.TaskId, err.Error()))
    }
    return int(okCnt), int(failedCnt)
}

// IsNeedTaskAlertDeal 判断是否需要执行alert分析处理
func IsNeedTaskAlertDeal(wpsName string) bool {
    wpsNames := strings.Split(wpsName, ",")
    wp := dbtask.Weapon{}
    wps, err := wp.SelectByNames(wpsNames)
    if err != nil {
        return true
    }

    for _, item := range wps {
        if item.Type == dbtask.WeaponTypeNeedAnalyzer {
            return true
        }
    }

    return false
}

/**
 等待任务结束
 */
func WaitTaskFinish(task *dbtask.WeaponTask, sumItem int, isNeedAnalyzerAlert bool) (chan int) {

    if sumItem <= 0 {
        return nil
    }

    notifyChan := make(chan int)

    go func(sum int) {
        doCount := 0
        for {
            select {
            case notify := <-notifyChan:
                if notify == 0 {
                    return
                }

                doCount = doCount + 1
                if doCount >= sum {
                    okMsg := fmt.Sprintf("task '%s' done over(weapon: %s; comment: %s)", task.TaskId, task.WeaponsName, task.Comment)
                    logs.Info(okMsg)
                    // 执行业务代码
                    task.OkCount, task.FailedCount = GetTaskCount(task)
                    task.Updated = time.Now().Unix()
                    task.Status = ctrutil.WEAPON_TASK_DO_OK
                    err := task.Update("status", "updated", "ok_count", "failed_count")
                    if err != nil {
                        errMsg := fmt.Sprintf("update task(%s)  failed: %s", task.TaskId, err.Error())
                        logs.Error(errMsg)
                        okMsg = okMsg + ";" + errMsg
                    }
                    // 结束
                    close(notifyChan)
                    // TODO: 通知类处理

                    // 执行数据生成处理
                    if isNeedAnalyzerAlert == true {
                        c := &analyzer.ComputeTaskAlerts{
                            TaskId: task.TaskId,
                        }
                        analyzer.ComputeQueueInstance().MsgQ.Push(c)
                    }
                    return
                }
            case <-time.After(time.Hour * 7 * 24):
                timeoutMsg := fmt.Sprintf("task '%s' happened timeout", task.TaskId)
                logs.Info(timeoutMsg)
                // 执行业务代码
                task.OkCount, task.FailedCount = GetTaskCount(task)
                task.Updated = time.Now().Unix()
                task.Status = ctrutil.WEAPON_TASK_DO_FAILED
                err := task.Update("status", "updated", "ok_count", "failed_count")
                if err != nil {
                    errMsg := fmt.Sprintf("update task(%s)  failed: %s", task.TaskId, err.Error())
                    logs.Error(errMsg)
                    timeoutMsg = timeoutMsg + ";" + errMsg
                }
                close(notifyChan)

                // TODO: 通知类处理
                return
            }
        }
    }(sumItem)

    return notifyChan
}

/**
 一个task item执行完成，发送Notify信息
 */
func (t *TaskItem) sendDoneNotify() error {
    var err error

    defer func() {
        if r := recover(); r != nil {
            err = fmt.Errorf("%v", r)
            logs.Error(fmt.Sprintf("task send notify: %s", err.Error()))
            return
        }
    }()

    t.NotifyChan <- 1
    return err
}

/**
 Done接口
 */
func (t *TaskItem) Done()  {

	// TaskItem必须要在结束之前执行，sendDoneNotify
    defer t.sendDoneNotify()

    // 找到对应的weapon
    wp := &dbtask.Weapon{Name: t.WeaponName}
    err := wp.Read("name")
    if err != nil {
        logs.Error(fmt.Sprintf("get weapon for '%s' error: %s", wp.Name, err.Error()))
        return
    }

    // 创建weapon log失败
    wpLog := t.insertData(wp)
    if wpLog == nil {
        return
    }

    // 发送weapon任务
    err = t.sendWeapon(wp)
    if err != nil {
        logs.Error(err.Error())
        wpLog.Updated = time.Now().Unix()
        wpLog.Status = ctrutil.WEAPON_LOG_FAILED
        wpLog.ErrorInfo = err.Error()
        err = wpLog.Update("updated", "status", "error_info")
        if err != nil {
            logs.Error(fmt.Sprintf("update weapon log failed info failed: %s", err.Error()))
        }
        return
    }

    // 下发成功
    wpLog.Updated = time.Now().Unix()
    wpLog.Status =  ctrutil.WEAPON_TASK_DOING
    err = wpLog.Update("updated", "status")
    if err != nil {
        logs.Error(fmt.Sprintf("update weapon log failed: %s", err.Error()))
    }
    return 
}

/**
 创建一行weapon log的数据
 */
func (t *TaskItem) insertData(wp *dbtask.Weapon) *dbdata.WeaponLog {
    if len(t.TaskId) != uconst.TaskIdLen {
        logs.Error(fmt.Sprintf("task item task id error"))
        return nil
    }
    if len(t.WormId) != uconst.AgentIdLen {
        logs.Error(fmt.Sprintf("task item worm id error"))
        return nil
    }

    wpLog := &dbdata.WeaponLog{
        WormId: t.WormId,
        WeaponId:    wp.Id,
        TaskId:      t.TaskId,
        Input:       t.Input,
        Created:     time.Now().Unix(),
        Status:      ctrutil.WEAPON_LOG_NO_DO,
    }
    wpLog.Updated = wpLog.Created
    err := wpLog.Insert()
    if err != nil {
        logs.Error(fmt.Sprintf("insert weapon log failed: %s", err.Error()))
        return nil
    }

    return wpLog
}

/**
 发送weapon的命令
 */
func (t *TaskItem) sendWeapon(wp *dbtask.Weapon) error {
    // 下发weapon之前检查agent正常在线
    wormId := &dbmanager.WormId{AgentId: t.WormId}
    err := wormId.Read("agent_id")
    if err != nil {
        return err
    }
    if wormId.Status != dbmanager.AgentOnline {
        return fmt.Errorf("agent not online: %s", wormId.IP)
    }

    // 发送出去
    args := urpc.WeaponMsgArgs{
        Id: t.WormId,
        TaskId: t.TaskId,
        WeaponId: int(wp.Id),
        WeaponName: wp.Name,
        WeaponMd5: wp.Md5,
        WeaponUrl: wp.Url,
        WeaponType: wp.Type,
        Timeout: wp.Timeout,
        Params: t.Input,
    }
    reply := urpc.WeaponMsgReply{}
    rpcAddr, _ := ctrutil.GetRpcIDForAccessSrv(wormId.ServerAddr)
    err = ctrutil.AccessSrvRpc.Call(rpcAddr, "ExecAction.ExecWp", &args, &reply)
    if err != nil {
        // update agent online status
        if err.Error() == "agent offline" {
            s := &slowtask.UpdateAgentOfflineStatus{
                AgentId: args.Id,
            }
            taskqueue.SlowQueueInst().ListC <- s
        }

        return fmt.Errorf("rpc call weapon for '%s' error: %s", args.Id, err.Error())
    }
    return nil
}

// @Description: 用于系统内部基于不同的业务动态创建weapon操作
type WeaponTask struct {
    TaskItem
    IP              string
    WaitResult      bool            // 同步等待结果
}

/**
 @Description：设置weapon task的参数，之前必须设置的参数有：ip、weapon_name、wait_result
 @Param:
 @Return：
 */
func (w *WeaponTask) setWeaponTaskParam() error {
    if net.ParseIP(w.IP) == nil {
        return fmt.Errorf("input ip error")
    }

    wormId := &dbmanager.WormId{}
    err := wormId.SelectValidByIP(w.IP)
    if err != nil {
        return fmt.Errorf("set weapon task error: %s", err.Error())
    }
    w.WormId = wormId.AgentId

    // 生成task id
    w.TaskId = uid.TaskIdProducer().CreateTaskId()

    return nil
}

/**
 @Description：动态下发的weapon结果
 @Param:
 @Return：
 */
func (w *WeaponTask) Done()  {

    err := w.setWeaponTaskParam()
    if err != nil {
        logs.Error(fmt.Sprintf("set weapon task param error: %s", err.Error()))
        return
    }

    if w.WaitResult == true {
        // 设置defer 执行通知
        // todo：xxxx
    }

    // 找到对应的weapon
    wp := &dbtask.Weapon{Name: w.WeaponName}
    err = wp.Read("name")
    if err != nil {
        logs.Error(fmt.Sprintf("get weapon for '%s' error: %s", wp.Name, err.Error()))
        return
    }

    // 创建weapon log失败
    wpLog := w.insertData(wp)
    if wpLog == nil {
        return
    }

    // 发送weapon任务
    err = w.sendWeapon(wp)
    if err != nil {
        logs.Error(err.Error())
        wpLog.Updated = time.Now().Unix()
        wpLog.Status = ctrutil.WEAPON_LOG_FAILED
        wpLog.ErrorInfo = err.Error()
        err = wpLog.Update("updated", "status", "error_info")
        if err != nil {
            logs.Error(fmt.Sprintf("update weapon log failed info failed: %s", err.Error()))
        }
        return
    }

    // 下发成功
    wpLog.Updated = time.Now().Unix()
    wpLog.Status = ctrutil.WEAPON_TASK_DOING
    err = wpLog.Update("updated", "status")
    if err != nil {
        logs.Error(fmt.Sprintf("update weapon log failed: %s", err.Error()))
    }
    return

}